# Deep Closest Point: Learning Representations for Point Cloud Registration
# https://github.com/WangYueFt/dcp/blob/bcd7deff0020de4be4db62cb953980d7b4f3a28c/model.py#L52

import torch
import torch.nn as nn
import math

class SVDHead(nn.Module):
    def __init__(self):
        super(SVDHead, self).__init__()
        self.reflect = nn.Parameter(torch.eye(3), requires_grad=False)
        self.reflect[2, 2] = -1

    def forward(self, *input):
        if len(input)==4:
            src_embedding = input[0]
            tgt_embedding = input[1]
            src = input[2]
            tgt = input[3]
            batch_size = src.size(0)

            d_k = src_embedding.size(1)
            scores = torch.matmul(src_embedding.transpose(2, 1).contiguous(), tgt_embedding) / math.sqrt(d_k)
            scores = torch.softmax(scores, dim=2)

            src_corr = torch.matmul(tgt, scores.transpose(2, 1).contiguous())
        else:
            src = input[0]
            tgt = input[1]
            batch_size = src.size(0)
            src_corr = tgt

        src_centered = src - src.mean(dim=2, keepdim=True)

        src_corr_centered = src_corr - src_corr.mean(dim=2, keepdim=True)

        H = torch.matmul(src_centered, src_corr_centered.transpose(2, 1).contiguous())

        U, S, V = [], [], []
        R = []

        for i in range(src.size(0)):
            u, s, v = torch.svd(H[i])
            r = torch.matmul(v, u.transpose(1, 0).contiguous())
            r_det = torch.det(r)
            if r_det < 0:
                u, s, v = torch.svd(H[i])
                v = torch.matmul(v, self.reflect)
                r = torch.matmul(v, u.transpose(1, 0).contiguous())
                # r = r * self.reflect
            R.append(r)

            U.append(u)
            S.append(s)
            V.append(v)

        U = torch.stack(U, dim=0)
        V = torch.stack(V, dim=0)
        S = torch.stack(S, dim=0)
        R = torch.stack(R, dim=0)

        t = torch.matmul(-R, src.mean(dim=2, keepdim=True)) + src_corr.mean(dim=2, keepdim=True)
        return R, t.view(batch_size, 3)


def cal_euler_from_Rmat(rot_mat):
    rx = torch.atan(rot_mat[2,1]/rot_mat[2,2])
    ry = torch.asin(-rot_mat[2,0])
    rz = torch.atan(rot_mat[1,0]/rot_mat[0,0])
    return torch.stack([-rx, -ry, rz])/torch.pi*180.0

if __name__ == '__main__':
    import numpy as np
    from wisp.core import transforms
    from scipy.spatial.transform import Rotation
    
    device = torch.device('cuda:1')
    T_mat = transforms.ObjectTransform(device=device)
    SVD_solver = SVDHead().to(device)

    pd_size = [8, 36, 3]
    src_pd = torch.rand(*pd_size, device=device)

    rot_array = []
    t_array = []
    mat_array = []
    tgt_pd_array = []
    for id in range(src_pd.shape[0]):
        # random transformation
        rot = torch.rand(3, device=device)*30
        t = torch.rand(3, device=device)*1000

        rot_array.append(rot.to('cpu').numpy())
        t_array.append(t.to('cpu').numpy())
        # rot += torch.randn_like(rot) # noise
        # t += torch.randn_like(t)

        T_mat.rx, T_mat.ry, T_mat.rz = rot[0], rot[1], rot[2]
        T_mat.tx, T_mat.ty, T_mat.tz = t[0], t[1], t[2]
        m = T_mat.model_matrix()
        mat_array.append(m)

        # rot_calc = cal_euler_from_Rmat(m)

        _src_pd = torch.cat([src_pd[id], torch.ones([src_pd.shape[1], 1], device=device)], dim=-1).transpose(0,1)
        _tgt_pd = (m @ _src_pd).transpose(0,1)
        tgt_pd_array.append(_tgt_pd[None, :, :-1])
        
    
    tgt_pd = torch.concat(tgt_pd_array, 0)
    mat_tensor = torch.stack(mat_array, dim=0)
    R_gt, t_gt = mat_tensor[:,:3, :3], mat_tensor[:, :3, -1]


    R_pred, t_pred = SVD_solver(src_pd.transpose(-2,-1), tgt_pd.transpose(-2,-1))

    # rotation degree absolute error
    # euler_array = []
    # for i in range(R_pred.shape[0]):
    #     r = Rotation.from_matrix(R_pred[i].detach().cpu().numpy())
    #     euler = r.as_euler('xyz', degrees=True)
    #     euler_array.append([-euler[0], -euler[1], euler[2]])
    # euler_np = np.stack(euler_array)
    
    euler_np = Rotation.from_matrix(R_pred.detach().cpu().numpy()).as_euler('xyz', degrees=True)
    euler_np[:, :2] = -euler_np[:, :2]

    rot_np = np.stack(rot_array)
    rot_deg_error = np.abs(euler_np-rot_np).sum(axis=-1).mean()

    # translation distance error
    transl_error = (t_pred-t_gt).norm(dim=-2).mean().to('cpu').numpy()
    t_np = np.stack(t_array)
    transl_error = np.linalg.norm(t_pred.to('cpu').numpy()-t_np, axis=-1).mean()

    print("Rot error: {:.6f}, t error: {:.6f}".format(rot_deg_error, transl_error))

    print("Done!")

